"""
Weeecore Bot 修复版本
根据嵌入式代码修复了I2C通信协议和返回值处理

作者: Weeecore Bot Team
版本: 1.0.3
"""

import time
import struct


class wm_bot:
    """
    Weeecore Bot 控制类
    
    通过I2C通信协议与下位机通信，实现各种机器人控制功能。
    支持多种控制模式：直接PWM控制、速度PID控制、位置PID控制等。
    
    根据嵌入式代码，I2C通信返回值规范：
    - 0x01: 指令执行成功
    - 0xFE: 指令执行失败
    """
    
    # 默认I2C从机地址
    BOT_ADDR = 0x28

    def __init__(self, i2c=None, addr=None):
        """
        初始化Weeecore Bot控制器
        
        Args:
            i2c: I2C对象，如果为None则尝试从全局变量获取
            addr: I2C从机地址，如果为None则使用默认地址0x28
            
        Raises:
            RuntimeError: 当板卡类型不支持或I2C对象未找到时
        """
        # 获取板卡类型，支持多种获取方式
        import builtins
        self.board_type = getattr(builtins, 'BoardType', None)

        # 如果builtins中没有，尝试从__main__模块获取
        if self.board_type is None:
            try:
                import __main__
                self.board_type = getattr(__main__, 'BoardType', None)
            except:
                pass

        # 验证板卡类型，确保是Weeecore Bot
        if self.board_type != "Weeecore Bot":
            raise RuntimeError(f"Unsupported board type: {self.board_type}")

        # 获取I2C对象，支持多种获取方式
        if i2c is None:
            try:
                import __main__
                if hasattr(__main__, 'i2c'):
                    self.i2c = __main__.i2c
                else:
                    raise RuntimeError("I2C object not found, please initialize I2C first")
            except:
                raise RuntimeError("I2C object not found, please initialize I2C first")
        else:
            self.i2c = i2c

        # 设置I2C从机地址
        self.bot_addr = addr or self.BOT_ADDR

        print(f"Weeecore Bot initialized successfully (address: 0x{self.bot_addr:02X})")

        # ==================== 私有方法 ====================

    def __send_and_read(self, send_data, read_bytes):
        """
        发送数据并读取指定字节数
        
        Args:
            send_data: 要发送的数据字节
            read_bytes: 要读取的字节数
            
        Returns:
            bytes: 读取到的数据，如果失败返回None
        """
        try:
            # 发送数据到I2C从机
            self.i2c.writeto(self.bot_addr, send_data)
            
            # 等待从机处理数据（5ms延时）
            time.sleep(0.005)
            
            # 如果需要读取数据，则从从机读取
            if read_bytes and read_bytes > 0:
                data = self.i2c.readfrom(self.bot_addr, read_bytes)
                return data
            return b""
        except Exception as e:
            print(f"I2C communication error: {e}")
            return None

    def __send_and_expect_ack(self, send_data):
        """
        发送数据并期望收到ACK应答
        
        Args:
            send_data: 要发送的数据字节
            
        Returns:
            bool: 是否成功收到ACK（0x01）
        """
        # 发送数据并读取1字节应答
        resp = self.__send_and_read(send_data, 1)
        
        # 检查应答是否有效
        if resp is None:
            print("I2C communication error: no response received")
            return False
        if len(resp) != 1:
            print(f"I2C communication error: response length abnormal ({len(resp)} bytes)")
            return False
        if resp[0] == 0xFE:
            print("I2C communication error: device returned error (0xFE)")
            return False
        elif resp[0] != 0x01:
            print(f"I2C communication error: response value abnormal (0x{resp[0]:02X}), expected 0x01")
            return False
        
        return True

    def __send_only(self, send_data):
        """仅发送数据，不读取应答"""
        try:
            self.i2c.writeto(self.bot_addr, send_data)
            return True
        except Exception as e:
            print(f"I2C send only error: {e}")
            return False

    def __bytes_to_ascii(self, data):
        """将字节数据转换为ASCII字符串"""
        if data is None:
            return ""
        ascii_string = ''.join(chr(byte) for byte in data if 32 <= byte <= 126)
        return ascii_string

    # ==================== 公有方法 ====================

    def get_bot_info(self):
        """获取设备基本信息 (0xAF)"""
        data = self.__send_and_read(b'\xAF', 48)
        if not data:
            return None
        
        # 查找成功标志0x01的位置，提取设备信息
        try:
            # 查找最后一个0x01作为成功标志
            for i in range(len(data)-1, -1, -1):
                if data[i] == 0x01:
                    # 提取设备信息（不包括成功标志）
                    info_data = data[:i]
                    return self.__bytes_to_ascii(info_data)
            # 如果没有找到成功标志，返回原始数据
            return self.__bytes_to_ascii(data)
        except Exception as e:
            print(f"Device info parsing failed: {e}")
            return None

    # ==================== LED控制指令 (0x10) ====================
    def led_on(self):
        result = self.__send_and_expect_ack(bytes([0x10, 0x01]))
        if not result:
            print("LED ON failed")
        return result

    def led_off(self):
        result = self.__send_and_expect_ack(bytes([0x10, 0x02]))
        if not result:
            print("LED OFF failed")
        return result

    def led_blink_disable(self):
        result = self.__send_and_expect_ack(bytes([0x10, 0x03]))
        if not result:
            print("LED blink disable failed")
        return result

    def led_blink_enable(self):
        result = self.__send_and_expect_ack(bytes([0x10, 0x04]))
        if not result:
            print("LED blink enable failed")
        return result

    # ==================== 电机控制指令 (0x20) ====================
    def motor_set_pwm(self, left_percent, right_percent):
        """直接PWM控制，输入范围-100..100 (整数)"""
        left = int(max(-100, min(100, left_percent))) & 0xFF
        right = int(max(-100, min(100, right_percent))) & 0xFF
        resp = self.__send_and_read(bytes([0x20, 0x01, left, right]), 1)
        return resp is not None and len(resp) == 1 and resp[0] == 0x01

    def motor_set_pwm_smooth(self, left_percent, right_percent, ramp_ms):
        """平滑PWM控制，ramp_ms按10ms量化并限幅0..2550ms"""
        left = int(max(-100, min(100, left_percent))) & 0xFF
        right = int(max(-100, min(100, right_percent))) & 0xFF
        ramp = int(max(0, min(255, ramp_ms // 10))) & 0xFF
        resp = self.__send_and_read(bytes([0x20, 0x02, left, right, ramp]), 1)
        return resp is not None and len(resp) == 1 and resp[0] == 0x01

    def motor_set_speed_rpm(self, left_rpm, right_rpm):
        """速度PID目标(-60..60 RPM)"""
        l = int(max(-60, min(60, left_rpm))) & 0xFF
        r = int(max(-60, min(60, right_rpm))) & 0xFF
        resp = self.__send_and_read(bytes([0x20, 0x03, l, r]), 1)
        return resp is not None and len(resp) == 1 and resp[0] == 0x01

    def motor_set_position(self, left_pos, right_pos):
        """位置PID目标，int32，按小端序发送"""
        # 修复：使用struct.pack代替int.to_bytes，避免参数问题
        lp = struct.pack('<i', int(left_pos))  # 小端序，有符号32位整数
        rp = struct.pack('<i', int(right_pos))  # 小端序，有符号32位整数
        resp = self.__send_and_read(b"\x20\x04" + lp + rp, 1)
        return resp is not None and len(resp) == 1 and resp[0] == 0x01

    def motor_stop_all(self):
        resp = self.__send_and_read(b"\x20\x05", 1)
        return resp is not None and len(resp) == 1 and resp[0] == 0x01

    def motor_get_status(self):
        # left_pos 左电机位置
        # right_pos 右电机位置
        # left_speed 左电机速度
        # right_speed 右电机速度
        # left_raw_speed 左电机原始速度
        # right_raw_speed 右电机原始速度
        # left_dir 左电机方向
        # right_dir 右电机方向
        """读取电机状态，返回字典: {left_pos,right_pos,left_speed,right_speed,left_raw_speed,right_raw_speed,left_dir,right_dir}"""
        data = self.__send_and_read(b"\x20\x06", 27)  # 26字节数据 + 1字节成功标志
        if not data or len(data) < 26:
            return None
        
        # 检查成功标志
        if len(data) == 27 and data[26] != 0x01:
            print(f"Motor status command failed, response: 0x{data[26]:02X}")
            return None
        
        try:
            # 按小端解析: 2个int32 + 4个float + 2个int8
            left_pos, right_pos, left_spd, right_spd, left_raw, right_raw, left_dir, right_dir = struct.unpack(
                '<iiffffbb', data[:26]
            )
            
            return {
                'left_pos': left_pos,
                'right_pos': right_pos,
                'left_speed': left_spd,
                'right_speed': right_spd,
                'left_raw_speed': left_raw,
                'right_raw_speed': right_raw,
                'left_dir': left_dir,
                'right_dir': right_dir,
            }
            
        except Exception as e:
            print(f"Motor status parsing failed: {e}")
            return None

    def motor_reset_position(self):
        resp = self.__send_and_read(b"\x20\x07", 1)
        return resp is not None and len(resp) == 1 and resp[0] == 0x01

    def motor_get_raw(self):
        """读取编码器原始数据，结构同motor_get_status"""
        data = self.__send_and_read(b"\x20\x08", 27)  # 26字节数据 + 1字节成功标志
        if not data or len(data) < 26:
            return None
        
        # 检查成功标志
        if len(data) == 27 and data[26] != 0x01:
            print(f"Motor raw data command failed, response: 0x{data[26]:02X}")
            return None
        
        try:
            left_pos, right_pos, left_spd, right_spd, left_raw, right_raw, left_dir, right_dir = struct.unpack(
                '<iiffffbb', data[:26]
            )
            
            return {
                'left_pos': left_pos,
                'right_pos': right_pos,
                'left_speed': left_spd,
                'right_speed': right_spd,
                'left_raw_speed': left_raw,
                'right_raw_speed': right_raw,
                'left_dir': left_dir,
                'right_dir': right_dir,
            }
            
        except Exception as e:
            print(f"Motor raw data parsing failed: {e}")
            return None

    def motor_set_speed_pid_params(self, kp, ki, kd):
        """设置速度PID参数（功能未实现）"""
        kp_bytes = struct.pack('<f', float(kp))
        ki_bytes = struct.pack('<f', float(ki))
        kd_bytes = struct.pack('<f', float(kd))
        resp = self.__send_and_read(b"\x20\x09" + kp_bytes + ki_bytes + kd_bytes, 1)
        return resp is not None and len(resp) == 1 and resp[0] == 0x01

    def motor_set_position_pid_params(self, kp, ki, kd):
        """设置位置PID参数（功能未实现）"""
        kp_bytes = struct.pack('<f', float(kp))
        ki_bytes = struct.pack('<f', float(ki))
        kd_bytes = struct.pack('<f', float(kd))
        resp = self.__send_and_read(b"\x20\x0A" + kp_bytes + ki_bytes + kd_bytes, 1)
        return resp is not None and len(resp) == 1 and resp[0] == 0x01

    # ==================== 电池电压读取 (0x30) ====================
    def get_battery_voltage(self):
        data = self.__send_and_read(b"\x30", 5)  # 4字节电压数据 + 1字节成功标志
        if not data or len(data) < 4:
            return None
        
        # 检查成功标志
        if len(data) == 5 and data[4] != 0x01:
            print(f"Battery voltage command failed, response: 0x{data[4]:02X}")
            return None
        
        try:
            voltage = struct.unpack('<f', data[:4])[0]
            if 0.0 <= voltage <= 20.0:
                return voltage
            else:
                return None
        except Exception as e:
            return None

    # ==================== 巡线传感器控制 (0x40) ====================
    def line_set_threshold(self, threshold):
        """设置阈值(0..65535)"""
        th = int(max(0, min(65535, threshold)))
        hb = (th >> 8) & 0xFF
        lb = th & 0xFF
        return self.__send_only(bytes([0x40, 0x01, hb, lb]))

    def line_read_values(self):
        """读取4路巡线传感器ADC，返回长度为4的列表"""
        data = self.__send_and_read(b"\x40\x02", 9)  # 8字节传感器数据 + 1字节成功标志
        if not data or len(data) < 8:
            return None
        
        # 检查成功标志
        if len(data) == 9 and data[8] != 0x01:
            print(f"Line sensor command failed, response: 0x{data[8]:02X}")
            return None
        
        try:
            v0, v1, v2, v3 = struct.unpack('<HHHH', data[:8])
            if all(0 <= v <= 4095 for v in [v0, v1, v2, v3]):
                return [v0, v1, v2, v3]
            else:
                return None
        except Exception as e:
            return None

    # ==================== 遥控器数据读取 (0x50) ====================
    def remote_read(self):
        """读取遥控器数据（功能未实现）"""
        data = self.__send_and_read(b"\x50", 1)
        if not data or len(data) != 1:
            return None
        
        success_flag = data[0]
        if success_flag == 0xFE:
            return None
        elif success_flag == 0x01:
            return True
        else:
            return None
